# %% 问题
"""
2022.2.1
1.很慢，运行一次180秒，书上是30秒，输的体无完肤。

2022.3.14
1.更慢了224s，但是在kaggle上白嫖GPU只需50s，而且环境不需要再次配置！！！

2022.6.17
1.改为超参数版本
"""

import sys
import time

# %%导包
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torchvision import transforms
from torchvision.datasets import ImageFolder


def evaluate_accuracy(data_iter, net, device=None):
    '''
    function:
        计算分类模型的准确率
    Parameters:

    '''
    if device is None and isinstance(net, torch.nn.Module):  # isinstance,net和torch.Module同一类型，则返回True
        device = list(net.parameters())[0].device
    acc_sum, n = 0, 0  # 初始化参数
    for x, y in data_iter:  # 通过这种方式直接接收参数
        if isinstance(net, torch.nn.Module):
            # 测试模型，会关闭dropout
            net.eval()
            """
            训练时是正对每个min-batch的，但是在测试中往往是针对单张图片，
            即不存在min-batch的概念。由于网络训练完毕后参数都是固定的，
            因此每个批次的均值和方差都是不变的，因此直接结算所有batch的均值和方差。
            所有Batch Normalization的训练和测试时的操作不同
            在训练中，每个隐层的神经元先乘概率P，然后在进行激活，在测试中，
            所有的神经元先进行激活，然后每个隐层神经元的输出乘P。
            即：没有batch分批操作，和dropout神经元失灵操作
            """
            # .cpu()保证可以数值加减,argmax:find(max())
            ''' 看不懂？？？'''
            acc_sum += (net(x.to(device)).argmax(dim=1) ==
                        y.to(device)).float().sum().cpu().item()
            # 改回训练模式
            net.train()
        else:
            acc_sum += (net(x.to(device)).argmax(dim=1) ==
                        y.to(device)).float().sum().cpu().item()
        n = n + y.shape[0]
    return acc_sum / n


# %%加载数据
# 超参数
batch_size = 32  # 每批里面的样本数
resize = 224  # 为什么变为224？？？

transform = transforms.Compose([
    transforms.Resize((resize, resize)),
    # transforms.RandomHorizontalFlip(p=0.5),
    # transforms.RandomRotation(25),
    # transforms.RandomVerticalFlip(p=0.1),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),

])

train_data = ImageFolder(r"E:\Class_exprience\Machine_Vision\9_class_design\training_set\training_set",
                         transform=transform)

test_data = ImageFolder(r"E:\Class_exprience\Machine_Vision\9_class_design\test_set\test_set", transform=transform)

# sys.platfrom 获取操作系统名称，这里是windows
# startswith 判断子字符串是否与主字符串一致
# 可省
if sys.platform.startswith('win'):
    # 0表示不用额外的进程来加速读取数据
    num_workers = 0
else:
    num_workers = 4

# torch.utils.data.DataLoader 小批量读取数据
# batch=235,batch_size=256,examples=60160 约为 60000
train_iter = torch.utils.data.DataLoader(
    train_data,  # 数据（特征，标签）
    batch_size=batch_size,  # 批量样本数
    shuffle=True,  # 打乱顺序，一般为True
    num_workers=num_workers)  # 额外进程，一般用于加速加载数据
test_iter = torch.utils.data.DataLoader(
    test_data,
    batch_size=batch_size,
    shuffle=False,
    num_workers=num_workers)

# %%建立模型
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.conv = nn.Sequential(
            # in_channels, out_channels, kernel_size, stride
            nn.Conv2d(3, 96, 11, 4),
            nn.ReLU(),
            # kernel_size, stride
            nn.MaxPool2d(3, 2),

            # 减小卷积窗口，使用步长为2来使得输入与输出的高和宽一致，
            # 且增大输出通道数
            nn.Conv2d(96, 256, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(3, 2),
            # 连续3个卷积层，且使用更小的卷积窗口。除了最后的卷积层外，
            # 进一步增大了输出通道数。
            # 前两个卷积层后不使用池化层来减小输入的高和宽
            nn.Conv2d(256, 384, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(384, 384, 3, 1, 1),
            nn.ReLU(),
            nn.Conv2d(384, 256, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(3, 2)
        )
        # 这里全连接层的输出个数比LeNet中的大数倍。使用丢弃层来缓解过拟合
        self.fc = nn.Sequential(
            nn.Linear(256 * 5 * 5, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(4096, 4096),
            nn.ReLU(),
            nn.Dropout(0.5),
            # 输出层。由于这里使用Fashion-MNIST
            # 所以用类别数为10，而非论文中的1000
            nn.Linear(4096, 2)
        )

    def forward(self, img):
        feature = self.conv(img)
        output = self.fc(feature.view(img.shape[0], -1))
        return output


net = AlexNet()

# 初始化模型参数，使之更容易收敛
for m in net.modules():
    if isinstance(m, (torch.nn.Linear, torch.nn.Conv1d, torch.nn.Conv2d)):
        torch.nn.init.xavier_normal_(m.weight)

# %%训练
# 超参数
lr, num_epochs = 0.1, 20  # lr为收敛步长，num_epochs为训练周期
optimizer = torch.optim.SGD(net.parameters(), lr=lr)  # 选择优化函数
loss = torch.nn.CrossEntropyLoss()  # 选择损失函数
writer = SummaryWriter('./log')  # 设置路径

# 将模型加载到指定运算器中
net = net.to(device)
n_train = len(train_data)
n_test = len(test_data)
print("training on ", device)
print("训练样本数：{} ， 测试样本数：{} ".format(n_train, n_test))

for epoch in range(num_epochs):
    print(15 * '*' + '第{}轮训练'.format(epoch) + 15 * '*')
    train_l_sum, train_acc_sum, n, batch_count = 0.0, 0.0, 0, 0
    start = time.time()

    # X(图片数据)和y(abel)是一批的数据，所以X是3维的 256*（28*28），y 256
    # 最后一批数据不满256，而是96的
    # y_hat=net(X),256*10
    for X, y in train_iter:
        X = X.to(device)
        y = y.to(device)
        y_pre = net(X)
        l = loss(y_pre, y)  # y_hat预测target概率，y真实target,l此批数据的loss
        # 梯度清零
        optimizer.zero_grad()
        l.backward()
        optimizer.step()  # 优化网络参数

        # 输出每批训练结果
        train_l_sum += l.cpu().item()
        train_acc_sum += (y_pre.argmax(dim=1) == y).sum().cpu().item()  # 计算正确样本数
        n += y.shape[0]  # n:累和每批样本数量，也就是训练样本数量
        batch_count += 1  # 批次
        if batch_count % 50 == 0:
            print('训练次数{} , 准确率：{} , loos:{} , time:{} '.format(batch_count, train_acc_sum / n_train, train_l_sum,
                                                                time.time() - start))
        writer.add_scalar('Train_loss', train_l_sum, batch_count + len(train_iter) * epoch)
        writer.add_scalar('Train_acc', train_acc_sum, batch_count + len(train_iter) * epoch)

    # 输出每轮测试结果
    test_acc = evaluate_accuracy(test_iter, net)
    print('测试acc:{} , time:{}'.format(test_acc, time.time() - start))
    writer.add_scalar('Test_acc', test_acc, epoch)
    torch.save(net.state_dict(), "model{}.pt".format(epoch))

# writer.add_graph(net, torch.rand(4, 3, resize, resize))
writer.close()

'''输出结果数据'''
y_test_pre = []
for x, y in test_iter:  # 通过这种方式直接接收参数
    net.eval()
    net = net.to(device)
    y_pre_matrix=net(x.to(device))
    y_pre = torch.argmax(y_pre_matrix, dim=1)
    y_test_pre.extend(y_pre.tolist())
y_test_pre = np.array(y_test_pre)
np.save('y_test_pre', y_test_pre)  # 保存为npy文件
